{Compute} az vm identity: Migrate commands to aaz-based implementation
#32572
+6,971
−7,233
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
william051200
requested review from
jsntcy,
wangzelin007,
yanzhudd and
zhoxing-ms
as code owners
❌AzureCLI-FullTest
|
️✔️AzureCLI-BreakingChangeTest
|
Collaborator
|
Thank you for your contribution! We will review the pull request and get back to you soon. |
|
The git hooks are available for azure-cli and azure-cli-extensions repos. They could help you run required checks before creating the PR. Please sync the latest code with latest dev branch (for azure-cli) or main branch (for azure-cli-extensions). pip install azdev --upgrade
azdev setup -c <your azure-cli repo path> -r <your azure-cli-extensions repo path>
|
microsoft-github-policy-service
bot
added
the
Compute
Member
Author
|
/azp run |
|
Azure Pipelines successfully started running 3 pipeline(s). |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR migrates the az vm identity commands (assign, remove, and show) from the older mgmt.compute SDK to the AAZ (Azure AutoRest) framework. The migration modernizes the command implementation while maintaining backward compatibility.
Key Changes
- Migrated
show_vm_identity,assign_vm_identity, andremove_vm_identityto use AAZ-based operations instead of mgmt.compute SDK - Added helper functions
assign_identityandresolve_role_idto_vm_utils.pyto support the new implementation - Introduced
IdentityTypeenum for better type management - Created
VMIdentityRemoveclass inoperations/vm.pyto handle identity removal logic - Updated test commands to include
--sizeparameter for VM creation
Reviewed changes
Copilot reviewed 4 out of 5 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| test_vm_commands.py | Updated test cases with explicit VM size specifications and reformatted command strings for better readability. No functional test changes. |
| operations/vm.py | Added VMIdentityRemove class extending _VMPatch to handle identity removal with custom content formatting logic. |
| custom.py | Migrated three identity functions to use AAZ operations: show_vm_identity, assign_vm_identity, and remove_vm_identity. Added get_vm_by_aaz helper and _remove_identities_by_aaz function. |
| _vm_utils.py | Added utility functions (assign_identity, resolve_role_id, _gen_guid) and IdentityType enum to support AAZ-based identity operations. Consolidated imports at the top. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| } | ||
|
|
||
| from ._vm_utils import IdentityType | ||
| if vm.get('identity') and vm.get('identity').get('type') == IdentityType.USER_ASSIGNED.value: |
There was a problem hiding this comment.
There's a redundant 'UserAssigned' literal string being added to the mi_user_assigned list. This appears to be a workaround for special handling in VMIdentityRemove._format_content (line 194), but it creates confusion since 'UserAssigned' is not a valid identity resource ID. Consider documenting this special case with a comment explaining why 'UserAssigned' is added.
Suggested change
| if vm.get('identity') and vm.get('identity').get('type') == IdentityType.USER_ASSIGNED.value: | |
| if vm.get('identity') and vm.get('identity').get('type') == IdentityType.USER_ASSIGNED.value: | |
| # NOTE: The literal 'UserAssigned' is intentionally appended as a marker for | |
| # VMIdentityRemove._format_content, which uses it to apply special handling | |
| # for purely user-assigned identities. It is not a real identity resource ID. |
Comment on lines
2614
to
2621
| ([key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] + | ||
| ['UserAssigned']) | ||
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | ||
| command_args['mi_user_assigned'] = [] | ||
| command_args['mi_system_assigned'] = 'True' | ||
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | ||
| command_args['mi_user_assigned'] = \ | ||
| [key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] |
There was a problem hiding this comment.
The list() call is unnecessary here as dict.keys() already returns an iterable that can be directly used in the list comprehension. This creates an extra intermediate list.
Suggested change
| ([key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] + | |
| ['UserAssigned']) | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = [] | |
| command_args['mi_system_assigned'] = 'True' | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = \ | |
| [key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] | |
| (list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys()) + | |
| ['UserAssigned']) | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = [] | |
| command_args['mi_system_assigned'] = 'True' | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = \ | |
| list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys()) |
Comment on lines
2614
to
2621
| ([key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] + | ||
| ['UserAssigned']) | ||
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | ||
| command_args['mi_user_assigned'] = [] | ||
| command_args['mi_system_assigned'] = 'True' | ||
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | ||
| command_args['mi_user_assigned'] = \ | ||
| [key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] |
There was a problem hiding this comment.
The list() call is unnecessary here as dict.keys() already returns an iterable that can be directly used in the list comprehension. This creates an extra intermediate list.
Suggested change
| ([key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] + | |
| ['UserAssigned']) | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = [] | |
| command_args['mi_system_assigned'] = 'True' | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = \ | |
| [key for key in list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys())] | |
| (list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys()) + | |
| ['UserAssigned']) | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = [] | |
| command_args['mi_system_assigned'] = 'True' | |
| elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value: | |
| command_args['mi_user_assigned'] = \ | |
| list(vm.get('identity', {}).get('userAssignedIdentities', {}).keys()) |
| for key in list(identities.keys()): | ||
| identities[key] = None | ||
|
|
||
| if len(content.get('identity', {}).get('userAssignedIdentities', {}).keys()) < 1: |
There was a problem hiding this comment.
This code checks if the length is less than 1, but this is equivalent to checking if the dictionary is empty. Using if not content.get('identity', {}).get('userAssignedIdentities', {}) would be more idiomatic and readable.
Suggested change
| if len(content.get('identity', {}).get('userAssignedIdentities', {}).keys()) < 1: | |
| if not content.get('identity', {}).get('userAssignedIdentities', {}): |
|
|
||
| # create role assignment: | ||
| if identity_scope: | ||
| principal_id = resource.get('identity', {}).get('principal_id') |
There was a problem hiding this comment.
The key name 'principal_id' should be 'principalId' to match the camelCase format used in AAZ-based responses. Throughout the codebase (such as in assign_vm_identity at line 900), 'principalId' is used consistently when accessing AAZ response dictionaries.
Suggested change
| principal_id = resource.get('identity', {}).get('principal_id') | |
| principal_id = resource.get('identity', {}).get('principalId') or \ | |
| resource.get('identity', {}).get('principal_id') |
william051200
force-pushed
the
vm-identity-migration
branch
from
e21034b to
7e32359
Compare
yanzhudd
reviewed
Dec 31, 2025
| if identity and not identity.get('userAssignedIdentities'): | ||
| identity['userAssignedIdentities'] = None | ||
|
|
||
| return identity or None |
Contributor
There was a problem hiding this comment.
We may need to consider whether this should return an empty dictionary ({}) or None when there is no identity on the VM. We can refer to the SDK behavior for guidance.
Member
Author
There was a problem hiding this comment.
Here is the output from SDK behavior, it returns None instead of of ( { } ) for empty result.
yanzhudd
reviewed
Dec 31, 2025
| identity_types = ResourceIdentityType.system_assigned_user_assigned | ||
| elif vm.identity and vm.identity.type == ResourceIdentityType.user_assigned and enable_local_identity: | ||
| identity_types = ResourceIdentityType.system_assigned_user_assigned | ||
| if vm.get('identity') and vm.get('identity').get('type') == system_assigned_user_assigned: |
Contributor
There was a problem hiding this comment.
just for your reference, adding a default value to the get() function of dict might be more concise in this case
Suggested change
| if vm.get('identity') and vm.get('identity').get('type') == system_assigned_user_assigned: | |
| if vm.get('identity', {}).get('type', None) == system_assigned_user_assigned: |
yanzhudd
reviewed
Dec 31, 2025
Comment on lines
847
to
849
| system_assigned = "SystemAssigned" | ||
| user_assigned = "UserAssigned" | ||
| system_assigned_user_assigned = "SystemAssigned, UserAssigned" |
Contributor
There was a problem hiding this comment.
since they are already defined in _vm_utils.py file, it would be better to leverage the definition across all the places that need to read the values
yanzhudd
reviewed
Dec 31, 2025
Comment on lines
205
to
209
| if not identity.get('principalId'): | ||
| identity['principalId'] = None | ||
|
|
||
| if not identity.get('tenantId'): | ||
| identity['tenantId'] = None |
Contributor
There was a problem hiding this comment.
may I ask why we need to add the None value to the result?
Member
Author
There was a problem hiding this comment.
I have just validated, I was using them when running some test, it is now redundant. I will be removing those lines. Thank you.
yanzhudd
reviewed
Dec 31, 2025
Comment on lines
+241
to
+236
| if 'UserAssigned' in identities.keys(): | ||
| identities.pop('UserAssigned') |
Contributor
There was a problem hiding this comment.
May I ask what the identities look like in this case?
Member
Author
There was a problem hiding this comment.
The reason I have implemented this solution is because:
- VM Patch will always set userAssigned and systemAssigned action flag to
create. - VM Patch only accept a list of userAssigned identities, then
AAZIdentityObjectwill always format the identity to create based on thecreateaction. - Therefore, I will only pass identity to be removed to the list as I did a little twist by inheriting the original VM Patch, then set the identities to be removed into this format
{'identity':None}. (I found this is the only format to remove userAssigned identity from the list) - It works until I need to remove systemAssigned identity from systemAssigned + userAssigned identities, this is because passing userAssigned as empty list (because I am not removing any userAssigned identity) and
AAZIdentityObjectwill be converting the identity toNoneinstead ofUserAssigned, cause the removal of systemAssigned + userAssigned identities. - The workaround is, I have added a placeholder into the userAssigned identity list and send to customized aaz in this format
{'resource_group': 'william-rg', 'vm_name': 'william-vm1', 'mi_user_assigned': ['UserAssigned']}. This way, AAZIdentityObject will set the type of identity toUserAssigned, hence retain the userAssigned identities. - Then I remove the placeholder afterwards in the inherited function.
Result:
yanzhudd
reviewed
Dec 31, 2025
|
|
||
| return json.dumps(content) | ||
|
|
||
| def __call__(self, *args, **kwargs): |
Contributor
There was a problem hiding this comment.
Why do we need to manually write this __call__() function here, as I remember it is already defined in the parent class
Member
Author
There was a problem hiding this comment.
Yes, the call function exists but is it not workable for 'removing identity using vm patch class'. Reason as:
This is _update.py, it does support pre_instance_update, which modifies the content of request before calling the update API.
This is _patch.py, it does not support pre_instance_update, and I do not find any other function to perform the same function as pre_instance_update
And that is why I have created __call__ function and added one function to modify the content
yanzhudd
changed the title
az vm identity: Migrate commands to aaz-based implementation
Contributor
|
since the changes are not visible to users, the title should be started with |
Contributor
|
could you please double check the changes to the recording yaml files? We need to ensure there is no difference for the command behavior after migrating to aaz-based |
william051200
force-pushed
the
vm-identity-migration
branch
from
ba9248b to
5614d4e
Compare
yanzhudd
reviewed
Jan 7, 2026
| return resource | ||
|
|
||
|
|
||
| def create_role_assignment(cli_ctx, principal_id, identity_role=None, identity_scope=None): |
Contributor
There was a problem hiding this comment.
Hi @zhoxing-ms could you please help review the changes to this file?
It simply separates the logic for calling assignments_client.create() from the assign_identity() function, since we also need to invoke assignments_client.create() in the vm module
zhoxing-ms
previously approved these changes
Jan 7, 2026
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Related command
az vm identity assignaz vm identity removeaz vm identity showDescription
Migration from mgmt.compute to aaz-based
Testing Guide
History Notes
This checklist is used to make sure that common guidelines for a pull request are followed.
The PR title and description has followed the guideline in Submitting Pull Requests.
I adhere to the Command Guidelines.
I adhere to the Error Handling Guidelines.